#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Time    : 2023/12/6 8:46
# @Author  : FuKai ZHANG
# @File    : yaml_control.py
# @Description :YAML格式文件的读写操作
import os
import yaml.scanner


class YamlDataOperation:
    """ 获取 yaml 文件中的数据 """

    def __init__(self, file_dir):
        self.file_dir = str(file_dir)

    def get_yaml_data(self) -> dict:
        """
        获取 yaml 中的数据
        :param: fileDir:
        :return:
        """
        # 判断文件是否存在
        if os.path.exists(self.file_dir):
            with open(self.file_dir, 'r', encoding='utf-8') as data:
                res = yaml.load(data, Loader=yaml.FullLoader)
        else:
            raise FileNotFoundError("文件路径不存在")
        return res

    def write_yaml_data(self, key: str, value) -> int:
        """
        更改 yaml 文件中的值, 并且保留注释内容
        :param key: 字典的key
        :param value: 写入的值
        :return:
        """
        with open(self.file_dir, 'r', encoding='utf-8') as file:
            lines = [line for line in file.readlines() if line != '\n']

        with open(self.file_dir, 'w', encoding='utf-8') as file:
            flag = 0
            for line in lines:
                left_str = line.split(":")[0]
                if key == left_str.lstrip() and '#' not in line:
                    newline = f"{left_str}: {value}"
                    line = newline
                    flag = 1
                file.write(f'{line}\n')
        return flag

